# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.base.book.util import pretty_book, human_readable_timestamp_record
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.bitstamp_v2.book_builder import BookBuilder
from coin.exchange.bitstamp_v2.kr_rest.product import BitstampProduct
from coin.exchange.bitstamp_v2.ws.private_subscriber import parse_channel_name


class BitstampHandler(WebsocketHandler):
  def __init__(self, *, event_dist=None, feed_checker_config=None):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)

  def _debug_print(self, bids, asks):
    self._logger.info('\nbids: %s\n asks: %s',
                      json.dumps(bids, indent=2),
                      json.dumps(asks, indent=2))

  def on_depth(self, record, queue_data, topic_data):
    data = record.data  # type: dict
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    if data.get('event') != 'data':
      self._logger.warning('Wrong book data received! %s', data)
      return

    native_symbol = parse_channel_name(data['channel'])['currency_pair']
    if not self.is_subscribed('book', native_symbol):
      return

    product = BitstampProduct.FromStrNativeProductNothrow(native_symbol)
    if product is None:
      self._logger.error('Book data comes earlier than subscription response!')
      return

    if native_symbol not in self._books:
      book = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
      self._books[native_symbol] = book
    else:
      book = self._books[native_symbol]
      book.clear_book()

    bids = data['data']['bids']
    asks = data['data']['asks']
    book.snapshot(bids, asks, record.timestamp, record.topic_id)
    self.ready = True

    if self._feed_checker:
      self._feed_checker.on_book(product=book.product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    native_symbol = parse_channel_name(data['channel'])['currency_pair']
    if not self.is_subscribed('trade', native_symbol):
      return

    product = BitstampProduct.FromStrNativeProductNothrow(native_symbol)
    if product is None:
      self._logger.error('Trade data comes earlier than subscription response!')
      return

    trade_info = data['data']
    price = float(trade_info['price'])
    qty = float(trade_info['amount'])
    if trade_info['type'] == 0:
      side = TradeSide.TRADE_BUY_SIDE
    elif trade_info['type'] == 1:
      side = TradeSide.TRADE_SELL_SIDE
    else:
      raise ValueError('Unknown side: %s', data)

    trade_event = create_market_trade_event(record.timestamp,
                                            native_symbol,
                                            price,
                                            qty,
                                            side,
                                            topic_id=record.topic_id)
    if self._feed_checker:
      self._feed_checker.on_trade(product=product,
                                  trade_event=trade_event,
                                  timestamp=record.timestamp)
    self.publish_event('trade', native_symbol, trade_event)
